from frames import ContoursFrame
from video import Video
import contours
from global_names import *
import random

training_bounds = {'canny1min':70, 'canny1max':110,'canny2min':200, 'canny2max':250, 'kern_min':1, 'kern_max':3}
n_train_frames = 20

class Trainer():
	"""
	This class is similar to Tracker, except it iterates through the video (actually randomly sampled sections of the video) many times, attempting the tracking process with varying random sets of parameters. After completing these iterations, it modifies the data file of the video to reflect the parameters that optimize tracking (currently defined as a minimal average change in contour size across frames).
	"""
	def __init__(self,video_dir,display=False,iterations=1000,log=True):
		self.video = Video(video_dir)
		self.display = display
		self.iterations = iterations
		self.log = log
	def train(self):
		avg_deltas = []
		params_all = []
		for trial in range(self.iterations):
			if self.log:
				print "%i/%i"%(trial+1,self.iterations)
				
			#training instance parameters
			self.temp_params = self.random_params(training_bounds, params_all)
			frame1 = random.randrange(n_train_frames,len(self.video.frames)) - n_train_frames
			frames = self.video.frames[frame1:frame1+n_train_frames]
		
			#frame iteration
			self.last_contour = None
			self.selected_contours = []
			self.deltas = []
			for frame in frames:
				self.process_frame(frame)
				
			avg_deltas.append(np.average(self.deltas))
			params_all.append(self.temp_params)
		
		#Conclude training:
		if len(avg_deltas) == 0:
			print "No good parameters found. Try more training iterations or close in tighter on cell."
		else:
			best_params = params_all[avg_deltas.index(min(avg_deltas))]
			self.store_data(best_params)
			
	def process_frame(self, frame):
		contours_frame = ContoursFrame(frame,self.temp_params)
		contour_idx, contour = contours.choose_contour(contours_frame.contours, self.last_contour)
		self.selected_contours.append(contour)
		delta = contours.delta_contour(contour,self.last_contour)
		if delta > 0:	
			self.deltas.append(delta)
		self.last_contour = contour
		
	def store_data(self, params):
		self.video.store_params(params)
		
	def random_params(self,training_bounds,used):
		params = []
		while params in used or params == []:
			canny1min = training_bounds['canny1min']
			canny1max = training_bounds['canny1max']
			canny2min = training_bounds['canny2min']
			canny2max = training_bounds['canny2max']
			kern_min = training_bounds['kern_min']
			kern_max = training_bounds['kern_max']
			params = {}
			params['brightness'] = random.choice(range(0,5))
			params['contrast'] = random.choice(range(1,3))
			params['canny_thresh1'] = random.choice(range(canny1min,canny1max+1))
			params['canny_thresh2'] = random.choice(range(canny2min,canny2max+1))
			params['kern_size'] = random.choice(range(kern_min,kern_max+1))
		return params